import os
import pandas as pd
import matplotlib.pyplot as plt
from collections import defaultdict

#import os
import pandas as pd
from collections import defaultdict
import matplotlib.pyplot as plt

# 配置
sizes = [1, 4, 8, 32]
num_batches = 5
metric_names = ["val_el10", "val_acc"]  # 移除 val_ma
colors = ['b', 'r']  # 只保留两个指标，两个颜色
output_dir = "outputs"
fig_save_dir = "metric_plots"
os.makedirs(fig_save_dir, exist_ok=True)

def read_public_before(size):
    path = os.path.join(output_dir, f"val_metrics_before_{size}_public.csv")
    return pd.read_csv(path)

def read_after(size, batch_idx):
    path = os.path.join(output_dir, f"val_metrics_after_{size}_{batch_idx}.csv")
    return pd.read_csv(path)

# ===== 任务 1：每个验证集单独画图 =====
for size in sizes:
    print(f"[INFO] Processing size = {size}")
    before_df = read_public_before(size)
    val_datasets = before_df["dataset"].tolist()
    
    for dataset in val_datasets:
        metrics_by_step = defaultdict(list)
        
        # step 0: before
        row = before_df[before_df["dataset"] == dataset]
        for metric in metric_names:
            metrics_by_step[metric].append(row[metric].values[0])

        # step 1-5: after
        for j in range(num_batches):
            after_df = read_after(size, j)
            row = after_df[after_df["dataset"] == dataset]
            if not row.empty:
                for metric in metric_names:
                    metrics_by_step[metric].append(row[metric].values[0])
            else:
                for metric in metric_names:
                    metrics_by_step[metric].append(None)

        steps = list(range(num_batches + 1))
        # 为每个 metric 分别画图
        for i, metric in enumerate(metric_names):
            plt.figure(figsize=(8, 4))
            plt.plot(steps, metrics_by_step[metric], marker='o', color=colors[i])
            if metric == "val_acc":
                plt.ylim(-0.02, 0.08)  # 可自行调整
            plt.title(f"Size {size} - {dataset} - {metric}")
            plt.xlabel("Training Step (0=before)")
            plt.tick_params(axis='y', labelleft=False)  # 取消 y 轴标签
            plt.grid(True)
            plt.tight_layout()
            fig_path = os.path.join(
                fig_save_dir, f"trend_size{size}_{dataset.replace('/', '_')}_{metric}.png"
            )
            plt.savefig(fig_path)
            plt.close()

# ===== 任务 2：各公共验证集指标均值变化图 =====
for size in sizes:
    avg_metrics = {metric: [] for metric in metric_names}

    # step 0: before
    before_df = read_public_before(size)
    for metric in metric_names:
        avg_metrics[metric].append(before_df[metric].mean())

    # step 1-5: after
    for j in range(num_batches):
        after_df = read_after(size, j)
        public_df = after_df.iloc[1:]  # 跳过第一行（当前训练集）
        for metric in metric_names:
            avg_metrics[metric].append(public_df[metric].mean())

    # 分别为两个指标画图
    steps = list(range(num_batches + 1))
    for i, metric in enumerate(metric_names):
        plt.figure(figsize=(8, 4))
        plt.plot(steps, avg_metrics[metric], marker='o', color=colors[i])
        if metric == "val_acc":
            plt.ylim(-0.02, 0.08)
        plt.title(f"{metric} mean across public val sets (size={size})")
        plt.xlabel("Training Step")
        plt.tick_params(axis='y', labelleft=False)  # 取消 y 轴标签
        plt.grid(True)
        plt.tight_layout()
        fig_path = os.path.join(fig_save_dir, f"avg_trend_size{size}_{metric}.png")
        plt.savefig(fig_path)
        plt.close()

print(f"[INFO] 所有图像已保存在：{fig_save_dir}/")
